import os
import numpy as np
import pandas as pd
from drone_model import DroneAeroModel, Config


# ====================== 配置区域 ======================
class MyConfig(Config):
    # 覆盖默认配置
    FEATURE_NAMES = ['acc', 'vel', 'gyro', 'pos']
    MODEL_SAVE_PATH = './trained_models/my_drone_model.pth'
    EPOCHS = 200
    K_SHOT = 32  # 适配样本数


# ====================== 增强版数据加载函数 ======================
def load_data_from_csv(file_path):
    """
    增强版数据加载函数，改进风况标签识别
    """
    df = pd.read_csv(file_path)

    # 数据校验
    required_columns = []
    for feat in MyConfig.FEATURE_NAMES:
        required_columns.extend([f"{feat}_{ax}" for ax in ['x', 'y', 'z']])
    for col in required_columns:
        if col not in df:
            raise ValueError(f"缺失必要列: {col}")

    # 加载特征数据
    data = {}
    for feat in MyConfig.FEATURE_NAMES:
        data[feat] = df[[f"{feat}_{ax}" for ax in ['x', 'y', 'z']]].values

    # 加载目标数据
    target_cols = [f"{MyConfig.TARGET_NAME}_{ax}" for ax in ['x', 'y', 'z']]
    data[MyConfig.TARGET_NAME] = df[target_cols].values

    # 改进的风况标签识别
    filename = os.path.basename(file_path).lower()
    if 'no_wind' in filename:
        data['condition'] = 'no_wind'
    elif 'light' in filename:
        data['condition'] = 'light_wind'
    elif 'moderate' in filename:
        data['condition'] = 'moderate_wind'
    elif 'strong' in filename:
        data['condition'] = 'strong_wind'
    elif 'gale' in filename:
        data['condition'] = 'gale_wind'
    else:
        data['condition'] = 'unknown'

    return data


# ====================== 主程序 ======================
def main():
    # 初始化配置和模型
    cfg = MyConfig()
    model = DroneAeroModel(cfg)

    # ===== 训练阶段 =====
    train_data_dir = "./data/train/"
    os.makedirs(train_data_dir, exist_ok=True)

    train_files = [
        os.path.join(train_data_dir, f)
        for f in os.listdir(train_data_dir)
        if f.endswith('.csv')
    ]

    if train_files:
        print(f"准备加载{len(train_files)}个训练文件...")
        raw_train_data = []
        for file in train_files:
            try:
                data = load_data_from_csv(file)
                raw_train_data.append(data)
                print(f"成功加载: {os.path.basename(file)} | 风况: {data['condition']}")
            except Exception as e:
                print(f"加载失败 {os.path.basename(file)}: {str(e)}")

        print("\n开始训练...")
        model.train(raw_train_data)
        model.save_model()
        print(f"训练完成，模型已保存到 {cfg.MODEL_SAVE_PATH}")
    else:
        print("警告: 未找到训练数据，跳过训练阶段")

    # ===== 预测阶段 =====
    test_data_dir = "./data/test/"
    os.makedirs(test_data_dir, exist_ok=True)

    test_files = [
        os.path.join(test_data_dir, f)
        for f in os.listdir(test_data_dir)
        if f.endswith('.csv')
    ]

    if not test_files:
        print("\n警告: 未找到测试数据，请检查路径: ./data/test/")
        return

    print("\n=== 开始预测 ===")
    for test_file in test_files:
        try:
            print(f"\n处理文件: {os.path.basename(test_file)}")
            test_data = load_data_from_csv(test_file)
            test_features = np.hstack([test_data[feat] for feat in cfg.FEATURE_NAMES])

            # 基础预测
            pred = model.predict(test_features)
            print(f"预测风况: {pred['wind_conditions'][0]}")
            print(f"气动力样本 (前3个):")
            for i in range(3):
                print(f"  {pred['forces'][i]}")

            # 带适配的预测
            if len(test_data[cfg.TARGET_NAME]) >= cfg.K_SHOT:
                X_adapt = np.hstack([test_data[feat][:cfg.K_SHOT] for feat in cfg.FEATURE_NAMES])
                Y_adapt = test_data[cfg.TARGET_NAME][:cfg.K_SHOT]

                adapted_pred = model.predict(test_features, (X_adapt, Y_adapt))
                print("\n带适配后的预测:")
                print(f"预测风况: {adapted_pred['wind_conditions'][0]}")
                print(f"气动力样本 (前3个):")
                for i in range(3):
                    print(f"  {adapted_pred['forces'][i]}")
            else:
                print("\n警告: 测试数据不足，无法进行适配预测")

        except Exception as e:
            print(f"处理文件 {os.path.basename(test_file)} 失败: {str(e)}")


if __name__ == "__main__":
    main()